import cv2
import numpy as np
from scipy.signal import find_peaks


def detect_shaxian__old(img):
    if isinstance(img, str):
        img = cv2.imread(img, cv2.IMREAD_COLOR)

    try:
        if img.shape[2] == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img
            img = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    except Exception as e:
        print(e)
    # gray = cv2.medianBlur(gray, 3)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    adaptive_image_mean = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 15, 7)
    gray = cv2.medianBlur(adaptive_image_mean, 3)
    # gray = cv2.GaussianBlur(adaptive_image_mean, (5, 5), 0)
    y1 = img.shape[0] // 2
    # y2 = adaptive_image_mean[y1, :]
    k = 100
    y2 = gray[y1 - k:y1 + k, :]
    y2 = np.median(y2, axis=0)
    y = 255 - y2
    peak_id, peak_property = find_peaks(y, height=5, prominence=30, distance=5)
    out = peak_id.tolist()
    return gray, out 


def detect_shaxian(img, pro):
    if isinstance(img, str):
        img = cv2.imread(img, cv2.IMREAD_COLOR)
    axx = np.all(img < 10)
    if axx:
        print('图片为黑色')
        return img, []
    else:
        im1, out = detect_shaxian__old(img.copy())

        a = len(out)
        if a == 0:
            im1, out = detect_shaxian__old(img)
            print("未检测到")

        elif a > int(pro):
            im1, out = detect_shaxian__old(img)
            print('个数超出')
        if len(out) == 0:
            im1 = im1
        elif len(out) > int(pro):
                im1 = im1
        else:
            im1 = draw_out(img.copy(), out)
    return im1, out


def draw_lines(im1, y1, ccc, k, color):
    if len(ccc) == 0:
        return im1
    ccc = np.array(ccc)
    im1[y1 - k:y1 + k, ccc] = color
    if 1:
        im1[y1 - k:y1 + k, ccc - 1] = color
        im1[y1 - k:y1 + k, ccc + 1] = color
    return im1


def draw_out(img, ccc):
    y1 = img.shape[0] // 2
    im1 = img.copy()
    colors = [(0, 255, 0), (0, 0, 255)]
    k = 100
    im1 = draw_lines(im1, y1, ccc, k, colors[0])
    ddd = []
    for i in range(len(ccc) - 1):
        d = ccc[i + 1] - ccc[i]
        if d > 13:
            ddd.append(ccc[i])
            ddd.append(ccc[i + 1])

    im1 = draw_lines(im1, y1, ddd, k, colors[1])
    cnt = len(ccc)
    cv2.putText(im1, f'{cnt}', (100, 100),
                cv2.FONT_HERSHEY_COMPLEX, 3, (0, 255, 255), 3)
    return im1


def test_detect_shaxian2(pa):
    import os
    import time
    pro = 399
    # fnlist = fnlist[]
    os.chdir(pa)
    try:
        os.mkdir('out')
    except:
        pass
    os.system('del out\\*.jpg')
    names = os.listdir(pa)
    for i in range(len(names)):
        fn = names[i]
        fn1, ext = os.path.splitext(fn)
        src = cv2.imread(fn, 1)

        if ext not in ['.jpg', '.jpeg', '.png', '.bmp']:
            continue

        # if 'Pic_2021_05_28_144213_13' not in fn1:
        #    continue
        # if 'Pic_2021_05_28_144000_6' not in fn1:
        #    continue
        # src = cv2.imdecode(np.fromfile(fn,dtype=np.uint8), 1)
        t0 = time.time()
        im, out = detect_shaxian(src, pro)
        t1 = time.time()
        print(f'{i} {len(out)} {fn1} {t1 - t0:.5f}')
        # cv2.waitKey(0)
        cv2.imwrite(f'out/{len(out)}_{fn1}_.jpg', im)

    cv2.destroyAllWindows()

def test_dedd(img):
    im1 = cv2.imread(img)
    if im1.shape[2] == 3:
        gray = cv2.cvtColor(im1, cv2.COLOR_BGR2GRAY)
    else:
        gray = im1
    # gray = cv2.medianBlur(gray, 3)
    gray = cv2.GaussianBlur(gray, (3, 3), 0)


    y1 = im1.shape[0] // 2
    # y2 = adaptive_image_mean[y1, :]
    k = 100
    y2 = gray[y1 - k:y1 + k, :]
    y2 = np.median(y2, axis=0)
    y = 255 - y2
    peak_id, peak_property = find_peaks(y, height=25, prominence=25, distance=3)
    out = peak_id.tolist()
    im1 = draw_out(im1.copy(), out)
    cv2.imwrite('out.jpg', im1)
    print(len(out))
if __name__ == '__main__':
    img = r'D:\data\save\shaxian\20240906\NG\20240906_133722'
    test_detect_shaxian2(img)
    # detect_shaxian__old(f'D:/data\save\test\71.jpg')
    # im11 = r'D:\data\save\shaxian\20240906\NG\20240906_120734\img_20240906_120734_0_ret.jpg'
    # test_dedd(im11)